import importlib
import importlib.util
import pickle
from functools import partial
from typing import *
from types import ModuleType
import os
import re
import pandas as pd

from config import ROOT_DIR

from utils.error import StrategyError


def get_factor(factors: List[str] = None) -> Dict[str, str]:
    path = os.path.join(ROOT_DIR, 'factors')

    if not os.path.exists(path):
        raise StrategyError(f'因子文件不存在')
    names: List[str] = os.listdir(path)

    fac: Dict[str, str] = {}
    for item in names:
        res = re.match(r'(\w*)\.py$', item)
        _p = os.path.join(path, item)
        if not os.path.isfile(_p) or res is None:
            continue
        if factors is not None and res.groups()[0] not in factors:
            continue
        with open(_p, 'r', encoding='utf-8') as f:
            codes: fac = f.read()
            fac[res.groups()[0]] = codes
    return fac


def get_strategies(strategies: List[str] = None) -> Dict[str, str]:
    path = os.path.join(ROOT_DIR, 'strategies')
    if not os.path.exists(path):
        raise StrategyError(f'策略文件不存在')
    names: List[str] = os.listdir(path)

    ste: Dict[str, str] = {}
    for item in names:
        res = re.match(r'(\w*)\.py$', item)
        _p = os.path.join(path, item)
        if not os.path.isfile(_p) or res is None:
            continue
        if strategies is not None and res.groups()[0] not in strategies:
            continue
        with open(_p, 'r', encoding='utf-8') as f:
            codes: ste = f.read()
            ste[res.groups()[0]] = codes
    return ste


def load_module(name: str, file_code: str) -> ModuleType:
    """
    字符串加载为module类型 并返回
    :param name:
    :param file_code:
    :return:
    """
    _mod: ModuleType = importlib.util.module_from_spec(
        importlib.util.spec_from_loader(name, loader=None)
    )
    try:
        exec(file_code, _mod.__dict__)
        return _mod
    except Exception as e:
        raise StrategyError('加载策略失败: %s - %s' % (name, e.__str__()))


class FactorLoader:
    def __init__(self, name: str, code: str = None):
        self.name = name
        self._mod = load_module(name, code if code is not None else get_factor([name])[name])

    def signal(self) -> Callable[[pd.DataFrame, Any, ...], pd.DataFrame]:
        """
        返回处理过的信号函数，将信号函数参数信息放入到闭包函数中
        :return:
        """
        func = getattr(self._mod, 'signal', None)
        if not callable(func):
            raise StrategyError(f'因子信号函数错误：{self.name}')

        def wrapper(df: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
            return func(df, *args, **kwargs)

        return wrapper


class StrategyLoader:
    def __init__(self, name: str, code: str = None):
        self.name = name
        self.arg_id: str = name
        self.mod = load_module(name, code if code is not None else get_strategies([name])[name])
        self.system_conf: Dict[str, Any] = getattr(self.mod, 'SYSTEM_CONF', {})

    @property
    def dataType(self) -> str:
        return self.system_conf.get('dataType', 'both')

    @property
    def withFundingRate(self) -> bool:
        return self.system_conf.get('withFundingRate', True)

    @property
    def factors(self) -> List[str]:
        """
        策略因子配置
        :return:
        """
        return self.system_conf.get('factors', [])

    @property
    def leverage(self) -> int:
        """
        杠杆
        :return:
        """
        return self.system_conf.get('leverage', 1)

    @property
    def positionSide(self) -> bool:
        """
        是否双向持仓
        :return:
        """
        return self.system_conf.get('positionSide', False)

    @property
    def factorFunc(self) -> Dict[str, Callable[[pd.DataFrame, Any, Any], pd.DataFrame]]:
        """
        因子信号函数
        :return:
        """
        if len(self.factors) == 0:
            return dict()
        return {k: FactorLoader(k, v).signal() for k, v in get_factor(self.factors).items()}

    @property
    def Kline(self):
        """
        策略基础K线配置
        :return:
        """
        return self.system_conf.get('Kline', 'hour')

    @property
    def symbols(self) -> List[str]:
        """
        策略选币池
        :return:
        """
        return self.system_conf.get('symbols', [])

    @property
    def historyKlineCount(self) -> int:
        """
        策略使用历史线数量
        :return:
        """
        return self.system_conf.get('historyKlineCount', 999)

    @property
    def KlineBatchSize(self) -> int:
        """
        策略K线合并数
        :return:
        """
        return self.system_conf.get('KlineBatchSize', 1)

    def is_execute(self) -> bool:
        func: Callable[[], bool] = getattr(self.mod, 'is_execute', None)
        if not callable(func):
            raise StrategyError(f'获取执行判断函数错误')
        return func()

    def injection(self, data: Dict[str, Any]):
        """
        注入策略计算所需数据
        :param data:
        :return:
        """
        system_data: type = getattr(self.mod, 'SystemData', None)
        if system_data is None or not isinstance(system_data, type):
            raise StrategyError(f'获取策略SystemData失败')
        try:
            for k, v in data.items():
                setattr(system_data, k, v)
        except Exception as e:
            raise StrategyError(f'注入数据失败：{e.__str__()}')

    def custom(self) -> bytes:
        """
        获取策略中自定义的缓存数据
        :return:
        """
        system_data: type = getattr(self.mod, 'SystemData', None)
        if system_data is None:
            raise StrategyError(f'获取策略SystemData失败')
        custom: Optional[Any] = getattr(system_data, 'custom', None)
        return pickle.dumps(custom)

    def pick(self, data: pd.DataFrame) -> pd.DataFrame:
        func: Callable[[pd.DataFrame], pd.DataFrame] = getattr(self.mod, 'pick', None)
        if not callable(func):
            raise StrategyError(f'获取策略选币函数失败')
        return func(data)

    def trade(self, data: pd.DataFrame):
        func: Callable[[pd.DataFrame], None] = getattr(self.mod, 'trade', None)
        if not callable(func):
            raise StrategyError(f'获取策略交易逻辑失败')
        return func(data)

    def backtest_data(self) -> Dict[str, pd.DataFrame]:
        system_data: type = getattr(self.mod, 'SystemData', None)
        if system_data is None:
            raise StrategyError(f'获取策略SystemData失败')
        data: Optional[Any] = getattr(system_data, 'backtest_data', None)
        setattr(system_data, 'backtest_data', {})
        if data is None:
            raise StrategyError(f'策略未定义回测数据')
        return data
